#include "groupByBatch.h"

namespace sp::Model::Group{

void registe_groupByBatch(){
    groupByBatch::registe();
}

groupByBatch::groupByBatch(std::vector<std::string> v): PythonFunction("/home/lx/SmartPipe/src/core/functions/Model/Group/", "", "groupByBatch", "groupByBatch", v){
    assert(v.size() == 1);
    name = "groupByBatch";
    class_name = "Model-Group-groupByBatch";
    type = GATHER;
    data_type_support.push_back(std::pair<std::vector<uint8_t>, uint8_t>({TENSOR},BATCH_TENSOR));
}

groupByBatch::groupByBatch(int batch_size): PythonFunction(){
    name = "groupByBatch";
    class_name = "Model-Group-groupByBatch";
    type = GATHER;
    pModulePath = "/home/lx/SmartPipe/src/core/functions/Model/Group/";
    pModuleHomePath = "";
    pModuleName = "groupByBatch";
    pClassName = "groupByBatch";
    this->batch_size = batch_size;
    data_type_support.push_back(std::pair<std::vector<uint8_t>, uint8_t>({TENSOR},BATCH_TENSOR));
}

groupByBatch::groupByBatch(): PythonFunction(){
    name = "groupByBatch";
    class_name = "Model-Group-groupByBatch";
    type = GATHER;
    pModulePath = "/home/lx/SmartPipe/src/core/functions/Model/Group/";
    pModuleHomePath = "";
    pModuleName = "groupByBatch";
    pClassName = "groupByBatch";
    data_type_support.push_back(std::pair<std::vector<uint8_t>, uint8_t>({TENSOR},BATCH_TENSOR));
}

groupByBatch::~groupByBatch(){

}

void groupByBatch::registe(){
    Function::Register("Model-Group-groupByBatch", createObject);
    Function::Register("Model-Group-groupByBatch", createObject2);
}

Function* groupByBatch::createObject(std::vector<std::string> params){
    return new groupByBatch(params);
}

Function* groupByBatch::createObject2(){
    return new groupByBatch();
}

void groupByBatch::start(){
    // 构建实例
    PythonFunction::defaultStart();
    args.push_back(std::to_string(batch_size));
    assert(args.size() == 1);
    pArgs = PyTuple_New(args.size());
    for(int i = 0; i < args.size(); i++){
        PyTuple_SetItem(pArgs, i, Py_BuildValue("s", args[i].c_str()));
    }
    pInstance = PyObject_CallObject(pConstruct, pArgs);
}

bool groupByBatch::waitForResource(){
    return true;
}

// DONE: 完善groupByBatch
/**
    vector<vector<Tensor*>> -> [Tensor,Tensor,...]
        groupByBatch: [Tensor,Tensor,...] -> [[],[],...,[Batch_Tensor],...]
    [[],[],...,[Batch_Tensor],...] -> vector<Batch_Tensor*>
**/
void groupByBatch::process(std::vector<std::vector<Data*>>& data_input, std::vector<Data*>& data_output){
    assert(data_input.size() != 0);
    assert(data_input[0].size() == 1);
    // 输入的tensor转成Python格式
    PyObject* data = convertToPython(data_input);
    // 调用python函数
    pArgs = PyTuple_New(1);
    PyTuple_SetItem(pArgs, 0, data);
    pReturn = PyObject_CallMethod(pInstance, "process", "O", pArgs);
    // 结果转换成cpp格式并添加到data_output
    convertToCpp(pReturn, data_input, data_output, 0);
    // 释放内存
    for(int i = 0; i < data_input.size(); i++){
        assert(data_input[i][0]->type == TENSOR);
        executor->freeAll(data_input[i][0]);
    }
}

bool groupByBatch::releaseResource(){
    return true;
}

void groupByBatch::finish(){
    // 把每组到batch的数据发送出来
    pReturn = PyObject_CallMethod(pInstance, "finish", "");
    std::vector<std::vector<Data*>> data_input;
    std::vector<Data*> data_output;
    convertToCpp(pReturn, data_input, data_output, 0);
    sendToQueues(data_output);
    PythonFunction::defaultFinish();
}

void groupByBatch::copy(Function* other){
    Function::defaultCopy(other);
    groupByBatch* oth = dynamic_cast<groupByBatch*>(other);
    this->batch_size = oth->batch_size;
}

};